import pickle
import os
from hashlib import sha256
import time
import rsa

from .leveldb_handler import LvlDB
from .nodes import *
from .encoding import *


### This Patricia Tree implementation does not use RLP-encoding, as I was interesting in logic pickle.dumps and pickle.loads could be swapped for rlp easily
class MerklePatriciaTree:
    debug = False

    def __init__(self, from_scratch=None, debug=None):
        if debug:
            self.debug = debug
        if from_scratch:
            self.db = LvlDB(debug=self.debug, from_scratch=True)
            self.root_hash = self.init_new_root()
        else:
            self.db = LvlDB(debug=self.debug, from_scratch=False)
            try:
                self.root_hash = self.db.get(b'root_hash')
            except KeyError:
                raise Exception("ERROR. Attempted to read \'root_hash\' from DB, but it was not found." + os.linesep +
                                "Repair your DB, delete it or request PatriciaTree to be built from scratch")

    def init_new_root(self):
        root = BranchNode()
        root_str = pickle.dumps(root)
        root_hash = binstr(sha256(root_str).hexdigest())
        self.db.put(root_hash, root_str)
        self.db.put(b'root_hash', root_hash)
        return root_hash

    def insert(self, key_nibbles, value): # key_nibbles have no prefix.
        curr_node_hash = self.root_hash
        curr_node = pickle.loads(self.db.get(curr_node_hash))
        curr_nibble_index = 0
        visited_hashes = list()
        while True:
            if self.debug: # Print the curr_node, curr_nibble_index, and the left key_nibbles
                f = open('./DEBUG.txt', 'a')
                print(curr_nibble_index, key_nibbles[curr_nibble_index:], file=f)
                if type(curr_node) == LeafNode:
                    print('This is a LeafNode', file=f)
                    print('Replace!', file=f)
                elif type(curr_node) == ExtensionNode:
                    print('This is a ExtensionNode', file=f)
                    print(binstr_to_nib(curr_node.key), file=f)
                elif type(curr_node) == BranchNode:
                    print('This is a BranchNode', file=f)
                    if not curr_node[key_nibbles[curr_nibble_index]]:
                        print('The corresponding nibble is NULL!', file=f)
                    else:
                        print('The corresponding nibble is NOT NULL!', file=f)
                print('\n', file=f)
            visited_hashes.append(curr_node_hash)
            ### CASE 0-2: The current node is a BranchNode
            if type(curr_node) == BranchNode:
                if not curr_node[key_nibbles[curr_nibble_index]]: # if hash_by_current_nibble = NULL
                    leaf_str = pickle.dumps(LeafNode(value))
                    leaf_hash = binstr(sha256(leaf_str).hexdigest())
                    self.db.put(leaf_hash, leaf_str)
                    if len(key_nibbles) - curr_nibble_index >= 2: # CASE 1
                        ## Create extension node pointing to new kv node, that stores the value.
                        extension = ExtensionNode(subnib_to_binstr(key_nibbles[curr_nibble_index + 1:]), leaf_hash)
                        extension_str = pickle.dumps(extension)
                        extension_hash = binstr(sha256(extension_str).hexdigest())
                        self.db.put(extension_hash, extension_str)
                        curr_node[key_nibbles[curr_nibble_index]] = extension_hash
                        curr_node_str = pickle.dumps(curr_node)
                        self.db.put(curr_node_hash, curr_node_str) # The curr_node is dirty, so _recursive_hash_update is needed.
                        self._recursive_hash_update(visited_hashes)
                        break
                    elif len(key_nibbles) - curr_nibble_index == 1: # CASE 0
                        # 	point branch directly to new kv node, that stores the value.
                        curr_node[key_nibbles[curr_nibble_index]] = leaf_hash
                        curr_node_str = pickle.dumps(curr_node)
                        self.db.put(curr_node_hash, curr_node_str)
                        self._recursive_hash_update(visited_hashes)
                        break
                    else:
                        raise Exception("len(key_nibbles)=", len(key_nibbles), "curr_nibble_index=", curr_nibble_index)
                else: # Else get down.
                    curr_node_hash = curr_node[key_nibbles[curr_nibble_index]]
                    curr_node_str = self.db.get(curr_node_hash)
                    curr_node = pickle.loads(curr_node_str)
                    curr_nibble_index += 1
            ### CASE 3-6: The current node is a ExtensionNode
            elif type(curr_node) == ExtensionNode:
                node_key = binstr_to_nib(curr_node.key)
                if node_key[0] == HEX_SEQUENCE_ODD:
                    node_key = node_key[1:]
                elif node_key[0] == HEX_SEQUENCE_EVEN:
                    node_key = node_key[2:]
                else:
                    raise ExtensionNode("Incorrect node_key:", node_key)
                ## CASE 3: The node_key covers the left key_nibbles
                if node_key == key_nibbles[curr_nibble_index:]:
                    # If ExtensionNode.key completely equal to rest of the key, it is definitely Extension, that points to Leaf, just overwrite!
                    new_child_str = pickle.dumps(LeafNode(value))
                    new_child_hash = binstr(sha256(new_child_str).hexdigest())
                    self.db.put(new_child_hash, new_child_str)
                    self.db.delete(curr_node.child_hash)
                    curr_node.child_hash = new_child_hash
                    curr_node_str = pickle.dumps(curr_node)
                    self.db.put(curr_node_hash, curr_node_str)
                    self._recursive_hash_update(visited_hashes)
                    break
                ## CASE 4-6: The node_key does not cover the left key_nibbles
                else:
                    ## Check how many nibbles are same
                    equal_elems = self._get_amount_of_equal_elements(node_key, key_nibbles[curr_nibble_index:])
                    if equal_elems == len(node_key): # CASE 4: The node_key equals to a part of the left key_nibbles
                        ## But it is not the whole remaining key, just keep on traversing
                        curr_node_hash = curr_node.child_hash
                        curr_node_str = self.db.get(curr_node_hash)
                        curr_node = pickle.loads(curr_node_str)
                        curr_nibble_index += equal_elems
                    elif equal_elems == 0: # CASE 5: The first nibble of the node_key does not equal to the first nibble of the left key_nibbles
                        ## If no equal elems, switch it to branch. Branch will have 2 entries: to old extension (or directly to leaf) and to new extention then to leaf (or directly to leaf).
                        branch = BranchNode()
                        leaf_str = pickle.dumps(LeafNode(value))
                        leaf_hash = binstr(sha256(leaf_str).hexdigest())
                        self.db.put(leaf_hash, leaf_str)
                        ## For the new extension (or directly leaf).
                        if len(key_nibbles[curr_nibble_index : ]) == 1: # Not to new extension, directly to new leaf.
                            branch[key_nibbles[curr_nibble_index]] = leaf_hash # Point directly to new leaf.
                        else: # To new extension, then to new leaf.
                            extension = ExtensionNode(subnib_to_binstr(key_nibbles[curr_nibble_index + 1:]), leaf_hash)
                            extension_str = pickle.dumps(extension)
                            extension_hash = binstr(sha256(extension_str).hexdigest())
                            self.db.put(extension_hash, extension_str)
                            branch[key_nibbles[curr_nibble_index]] = extension_hash # Point to new extension.
                        ## For the old extension
                        self.db.delete(curr_node_hash)
                        if len(node_key) == 1: # Not to old extension, directly to old leaf.
                            branch[node_key[0]] = curr_node.child_hash # Point directly to old leaf.
                        else: # To old extension (node_key changes), then to old leaf.
                            curr_node.key = subnib_to_binstr(node_key[1:])
                            curr_node_str = pickle.dumps(curr_node)
                            curr_node_hash = binstr(sha256(curr_node_str).hexdigest())
                            branch[node_key[0]] = curr_node_hash # Point to old but changed extension.
                            self.db.put(curr_node_hash, curr_node_str)
                        branch_str = pickle.dumps(branch)
                        branch_hash = binstr(sha256(branch_str).hexdigest())
                        self.db.put(branch_hash, branch_str)
                        old_curr_node_hash = visited_hashes.pop() ## Since self.db.delete(curr_node_hash)
                        prev_node_hash = visited_hashes[-1]
                        prev_node = pickle.loads(self.db.get(prev_node_hash))
                        if type(prev_node) != BranchNode:
                            raise Exception("prev_node is of unexpected type:", type(curr_node))
                        ## The contents of prev_node (a branch) changes, since only branch can point to extension
                        nibble = prev_node.get_index(old_curr_node_hash)
                        prev_node[nibble] = branch_hash
                        prev_node_str = pickle.dumps(prev_node)
                        self.db.put(prev_node_hash, prev_node_str)
                        self._recursive_hash_update(visited_hashes)
                        break
                    ## CASE 6: A part of the node_key equals to a part of the left key_nibbles
                    elif equal_elems < len(node_key): # Have to split into extension -> branch -> ???.
                        leaf_str = pickle.dumps(LeafNode(value))
                        leaf_hash = binstr(sha256(leaf_str).hexdigest())
                        branch = BranchNode()
                        curr_nibble_index += equal_elems
                        self.db.put(leaf_hash, leaf_str)
                        ## CASE 6_1: Barring the equivalent nibbles, there is one nibble left in key_nibbles
                        if len(key_nibbles[curr_nibble_index:]) == 1: # Directly to a new leaf (extension -> branch -> leaves).
                            branch[key_nibbles[-1]] = leaf_hash
                            branch[node_key[equal_elems]] = curr_node.child_hash
                        ## CASE 6_2: Barring the equivalent nibbles, there is more than one nibble left in key_nibbles
                        else: # Extension -> branch -> extensions -> leaves
                            ## extension_to_leaf is the extension to new leaf
                            extension_to_leaf = ExtensionNode(subnib_to_binstr(key_nibbles[curr_nibble_index + 1:]), leaf_hash)
                            extension_to_leaf_str = pickle.dumps(extension_to_leaf)
                            extension_to_leaf_hash = binstr(sha256(extension_to_leaf_str).hexdigest())
                            self.db.put(extension_to_leaf_hash, extension_to_leaf_str)
                            branch[key_nibbles[curr_nibble_index]] = extension_to_leaf_hash
                            ## ext_after_branch is the extension to old leaf
                            ext_after_branch = ExtensionNode(subnib_to_binstr(node_key[equal_elems + 1:]), curr_node.child_hash)
                            ext_after_branch_str = pickle.dumps(ext_after_branch)
                            ext_after_branch_hash = binstr(sha256(ext_after_branch_str).hexdigest())
                            self.db.put(ext_after_branch_hash, ext_after_branch_str)
                            branch[node_key[equal_elems]] = ext_after_branch_hash
                        branch_str = pickle.dumps(branch)
                        branch_hash = binstr(sha256(branch_str).hexdigest())
                        self.db.put(branch_hash, branch_str)
                        ## The contents of curr_node changes, since this extension goes to new branch
                        ext_before_branch = ExtensionNode(subnib_to_binstr(node_key[:equal_elems]), branch_hash)
                        ext_before_branch_str = pickle.dumps(ext_before_branch)
                        ext_before_branch_hash = binstr(sha256(ext_before_branch_str).hexdigest())
                        self.db.delete(curr_node_hash)
                        self.db.put(ext_before_branch_hash, ext_before_branch_str)
                        visited_hashes.pop() # Since self.db.delete(curr_node_hash)
                        prev_node_hash = visited_hashes[-1]
                        prev_node = pickle.loads(self.db.get(prev_node_hash))
                        ## The contents of prev_node (a branch) changes, since only branch can point to extension
                        nibble = prev_node.get_index(curr_node_hash)
                        prev_node[nibble] = ext_before_branch_hash
                        prev_node_str = pickle.dumps(prev_node)
                        self.db.put(prev_node_hash, prev_node_str)
                        self._recursive_hash_update(visited_hashes)
                        break
                    elif equal_elems > len(node_key):
                        raise Exception("equal_elems(", equal_elems, ") > len(node_key) (", len(node_key), ")")
            elif type(curr_node) == LeafNode:
                raise Exception("curr_node is LeafNode!")
            else:
                raise Exception("curr_node is of unknown type:", type(curr_node))

    def remove(self, key_nibbles):
        curr_node_hash = self.root_hash
        curr_node = pickle.loads(self.db.get(curr_node_hash))
        curr_nibble_index = 0
        visited_hashes = list()
        while True:
            if self.debug: # Print the curr_node, curr_nibble_index, and the left key_nibbles
                f = open('./DEBUG.txt', 'a')
                print(curr_nibble_index, key_nibbles[curr_nibble_index:], file=f)
                if type(curr_node) == ExtensionNode:
                    print('This is a ExtensionNode', file=f)
                    print(binstr_to_nib(curr_node.key), file=f)
                elif type(curr_node) == BranchNode:
                    print('This is a BranchNode', file=f)
                print('\n', file=f)
            visited_hashes.append(curr_node_hash)
            if type(curr_node) == BranchNode:
                if not curr_node[key_nibbles[curr_nibble_index]]: # hash_by_current_nibble = NULL:
                    print('The node to be removed does not exist!')
                elif len(key_nibbles) == curr_nibble_index + 1: # key_nibbles ends, erase the nibble of the branch.
                    self.db.delete(curr_node[key_nibbles[curr_nibble_index]]) # Delete the leaf node.
                    curr_node[key_nibbles[curr_nibble_index]] = b'' # I do not consider whether the branch node is empty or not.
                    curr_node_str = pickle.dumps(curr_node)
                    self.db.put(curr_node_hash, curr_node_str)
                    self._recursive_hash_update(visited_hashes)
                    return
                else: # Get down
                    curr_node_hash = curr_node[key_nibbles[curr_nibble_index]]
                    curr_node_str = self.db.get(curr_node_hash)
                    curr_node = pickle.loads(curr_node_str)
                    curr_nibble_index += 1
            elif type(curr_node) == ExtensionNode:
                node_key = binstr_to_subnib(curr_node.key)
                equal_elems = self._get_amount_of_equal_elements(node_key, key_nibbles[curr_nibble_index:])
                if equal_elems != len(node_key):
                    print('The node to be removed does not exist!')
                elif node_key == key_nibbles[curr_nibble_index:]: # The extension node is the last node before the leaf node.
                    self.db.delete(curr_node.child_hash) # Delete the leaf node.
                    self.db.delete(curr_node_hash) # Delete the extension node.
                    visited_hashes.pop()
                    ## Deal with the previous branch node.
                    prev_node_hash = visited_hashes[-1]
                    prev_node = pickle.loads(self.db.get(prev_node_hash))
                    nibble = prev_node.get_index(curr_node_hash)
                    prev_node[nibble] = b''
                    prev_node_str = pickle.dumps(prev_node)
                    self.db.put(prev_node_hash, prev_node_str)
                    self._recursive_hash_update(visited_hashes)
                    return
                else: # Get down
                    curr_node_hash = curr_node.child_hash
                    curr_node_str = self.db.get(curr_node_hash)
                    curr_node = pickle.loads(curr_node_str)
                    curr_nibble_index += len(node_key)


    @staticmethod
    def _get_amount_of_equal_elements(array1, array2):
        """
        :return: amount of equal elements
        """
        for i in range(min(len(array1), len(array2))):
            if array1[i] != array2[i]:
                return i
        return min(len(array1), len(array2))

    def _recursive_hash_update(self, hashes):
        """type hashes: list

        Usage: top node in hashes should be left inconsistent, e.g. its key(hash)
        does not correspond to actual hash"""
        ### Handle top node, which only has hash inconsistency, as opposed to other nodes, which have to update links

        prev_node_former_hash = hashes.pop() # Removes the last element and returns the value of that element.
        prev_node_str = self.db.get(prev_node_former_hash)
        prev_node_real_hash = binstr(sha256(prev_node_str).hexdigest())
        self.db.delete(prev_node_former_hash)
        self.db.put(prev_node_real_hash, prev_node_str)

        while len(hashes):
            curr_hash = hashes.pop()
            curr_node_str = self.db.get(curr_hash)
            curr_node = pickle.loads(curr_node_str)
            if type(curr_node) == BranchNode:
                nibble = curr_node.get_index(prev_node_former_hash)
                curr_node[nibble] = prev_node_real_hash
            elif type(curr_node) == ExtensionNode:
                curr_node.child_hash = prev_node_real_hash
            elif type(curr_node) == LeafNode:
                raise Exception("Recursive hash update hit LeafNode")
            else:
                raise Exception("Incorrect type:", type(curr_node))
            prev_node_former_hash = curr_hash
            prev_node_str = pickle.dumps(curr_node)
            prev_node_real_hash = binstr(sha256(prev_node_str).hexdigest())

            if prev_node_former_hash != prev_node_real_hash:
                self.db.delete(prev_node_former_hash)
                self.db.put(prev_node_real_hash, prev_node_str)

        # Update root node
        self.root_hash = prev_node_real_hash
        self.db.put(b'root_hash', self.root_hash)
